'''这是实现dmrg二维参数扫描的程序
首先让dmrg计算一组二维的参数
然后把pyalps带的collectXY抓取的DataSet中的信息提出列成相图一样的二维矩阵，并附上x和y，这就是二维扫描参数要的最好的结果了
'''

# encoding: utf-8
# if you are using spyder, run these:
# from IPython import get_ipython
# get_ipython().magic('reset -sf')
 
import sys
sys.path.append('/opt/alps/lib')
import pyalps
import matplotlib.pyplot as plt
import pyalps.plot
import numpy as np
import time
import scipy.io
 
#----------------------1 Run the Task-------------------
a1 = time.time()
BHparms = []  ## 用dictionary作为元素组成list
Js = np.arange(2, 2.2, 0.04)
mus = np.arange(0, 2.2, 0.4)
for J in Js:
   for mu in mus:
       BHparms.append({
           'LATTICE': "open chain lattice",
           'L': 4,
           'MODEL_LIBRARY': "bhmodel.xml",
           'MODEL': "my boson hubbard",
           'J': J,
           'U': 1,
           'mu': mu,
           'Nmax': 6,
           'N_total': 8,
           'MAXSTATES': 128,
           'SWEEPS': 10,
           'MEASURE_LOCAL[expectation_n]': 'n(i)',
           'MEASURE_LOCAL[expectation_b]': 'b(i)'
           #        'MEASURE_LOCAL[expectation_n^2]':'n(i)*n(i)'
           #        'MEASURE_LOCAL[expectation_b^2]':'b(i)*b(i)'
       })
print('report: This 2D sweep task has %d parallel tasks.' % (len(BHparms)))
print('Running...')
input_file = pyalps.writeInputFiles('2Dsweep', BHparms)
pyalps.runApplication('dmrg', input_file, writexml=True, MPI=3)
a2 = time.time()
print('rundmrg has spent %.1f seconds' % (a2 - a1))
 
#----------------------2 结果分析--------------------------
## 注意这些和上面运行task部分的要保持一致
Js = np.arange(2, 2.2, 0.04)
mus = np.arange(0, 2.2, 0.4)
## 总共几个格点和后面数据处理选中间格点时还有关，请随之修改
L = 4
 
result_files = pyalps.getResultFiles(prefix='2Dsweep')
 
## 期望值b
eigen_measure_b = pyalps.loadEigenstateMeasurements(result_files,
                                                    'expectation_b')
bs = pyalps.collectXY(eigen_measure_b, 'J', 'expectation_b', ['mu'])
#print(bs)
def by_mu(dataset):
    return dataset.props['mu']
bs_sorted = sorted(bs, key=by_mu)
J_mu_2D_b = []
for i in bs_sorted:
    middle 
_values = [j[1] for j in i.y]  ## 这是想索引到中间格点的期望值
    J_mu_2D_b.insert(0, middle_values)  ## 低mu的最后在下面
J_mu_2D_b = np.array(J_mu_2D_b)
 
print('J_mu_2D_b is \n' + str(J_mu_2D_b))
print('Js is ' + str(Js))
print('mus is ' + str(mus))
 
## 期望值n
eigen_measure_n = pyalps.loadEigenstateMeasurements(result_files,
                                                    'expectation_n')
#print(eigen_measure_n)
ns = pyalps.collectXY(eigen_measure_n, 'J', 'expectation_n', ['mu'])
 
#print(ns)
## nice,collectXY产生的是最容易处理数据的了
## 单个J对应一些mu，有很多J，但是J在list中是乱序的
## 我打算把它排序
## b(i)，n(i)对于BH model只需要中间格点期望值
def by_mu(dataset):
    return dataset.props['mu']
 
ns_sorted = sorted(ns, key=by_mu)
## 把DataSet按照其props中的mu值排序，可以实现使用其中1个数据将DataSet这种自定义数据排序
#print(ns_sorted)
J_mu_2D_n = []
for i in ns_sorted:
middle_values = [j[1] for j in i.y]  
## 这是想索引到中间格点的期望值
J_mu_2D_n.insert(0, middle_values)  
## list的insert，插入数据到第一位，低mu的最后在下面
J_mu_2D_n = np.array(J_mu_2D_n)
 ## 最好的数据形式了：
print('J_mu_2D_n is \n' + str(J_mu_2D_n))
print('Js is ' + str(Js))
print('mus is ' + str(mus))
 
## 然后可以将结果保存为matlab的mat文件吧，这样能在Python和matlab之间实现数据互通
scipy.io.savemat('test_2233.mat', {
    'J_mu_2D_b': J_mu_2D_b,
    'J_mu_2D_n': J_mu_2D_n,
    'Js': Js,
    'mus': mus
})
 
## 画图先不谈了
#fig=plt.figure()
#(fig,ax)=plt.subplots(figsize=(4,3),dpi=100)
#ax.plot(x1,y1,'b',label='|Magneti|')
#ax.plot(x2,y2,'--',color='purple',label='Energy')
#ax.legend(loc=4)
#ax.set_xlabel('x')
#ax.set_ylabel('y')
#ax.set_title('test tu')
#fig.savefig('dema.pdf')

'''
具体运行结果为：
In [4]: runfile('/home/aa/Desktop/alps数据处理/Brute_force_2D_xiangtu/BHmodel_dmrg_2Dsweep.py', wdir='/home/aa/Desktop/alps数据处理/Brute_force_2D_xiangtu')
J_mu_2D_b is 
[[  1.41937030e-05   1.25428619e-05  -1.37717307e-05   2.04229621e-06
    1.54585978e-05   3.69577071e-06]
 [ -1.03945051e-05  -3.94049781e-06   2.14163268e-06   1.74123849e-05
    2.79125632e-05  -7.21025955e-06]
 [  2.13842308e-06   1.36867827e-06   1.37865874e-05   1.24496517e-05
   -7.52171254e-06  -1.56340307e-06]
 [  1.29480667e-05  -2.92618610e-05  -2.08710467e-05   2.64308374e-05
    5.39490502e-06  -1.09435589e-05]
 [  2.63502920e-05   6.91557931e-06  -5.84492883e-06  -6.57038771e-06
    1.51401311e-06  -1.08105390e-05]
 [ -1.51366850e-05   4.06324689e-06   6.98041379e-06  -3.63688615e-07
    5.79743726e-06   9.22373607e-07]]
Js is [ 2.    2.04  2.08  2.12  2.16  2.2 ]
mus is [ 0.   0.4  0.8  1.2  1.6  2. ]


J_mu_2D_n is 
[[ 4.12246258  4.12273832  4.12298943  4.123229    4.12346171  4.12370755]
 [ 3.96533581  3.96625079  3.96713857  4.12322512  4.12345305  4.12368011]
 [ 3.96532723  3.96627879  3.96715129  3.9680052   3.96886339  3.9696539 ]
 [ 3.80485512  3.80662709  3.80831529  3.80994323  3.8115474   3.8131074 ]
 [ 3.63374334  3.6363737   3.63897196  3.64147896  3.6439226   3.64630991]
 [ 3.44611362  3.44969801  3.45316588  3.45656709  3.45984638  3.64634254]]
Js is [ 2.    2.04  2.08  2.12  2.16  2.2 ]
mus is [ 0.   0.4  0.8  1.2  1.6  2. ]

这就是实现二维参数扫描的全套代码了
'''